import requests
from gmssl.sm4 import CryptSM4, SM4_ENCRYPT
import base64
# pip install gmssl -i https://pypi.tuna.tsinghua.edu.cn/simple


# 密码加密函数
def mima_SM4(password):
    # 密钥
    SM4_key = 'tiekeyuankp12306'
    # 创建SM4加密对象
    crypt_sm4 = CryptSM4()
    # 设置密钥
    crypt_sm4.set_key(SM4_key.encode('utf-8'), SM4_ENCRYPT)
    # 加密密码
    encrypted_password = crypt_sm4.crypt_ecb(password.encode('utf-8'))
    # 将加密后的密码转换为base64编码的字符串
    encrypted_password_base64 = base64.b64encode(encrypted_password).decode('utf-8')
    encrypted_password_base64 = '@' + encrypted_password_base64
    print('加密后的密码:',encrypted_password_base64)
    return encrypted_password_base64

# 登入函数
'''
传参：
    username:账号
    castNum:身份证后4位
    password:密码
返回值：
    session:登入后的会话对象
'''
def login(username,castNum,password):
    # 设置请求头
    headers_1 = {
        'User-Agent':'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36',
    }

    # 创建会话对象
    session = requests.Session()
    session.get('https://www.12306.cn/index/')
    session.get('https://kyfw.12306.cn/otn/resources/login.html')
    # 构造表单数据
    data_1 = {
        'username':username,                 # 账号,手机号
        'appid':'otn',
    }
    data_2 = {
        'appid':'otn',
        'username':username,                 # 账号
        'castNum':castNum                    # 身份证后6位
    }
    url_login = 'https://kyfw.12306.cn/passport/web/checkLoginVerify'
    url_yanzheng = 'https://kyfw.12306.cn/passport/web/getMessageCode'

    # 进行登入
    # 获取验证码
    response = session.post(url_login,data=data_1,headers=headers_1)               # 登入检查
    print(response.json())
    response = session.post(url_yanzheng,data=data_2)
    if '成功' in response.json()['result_message']:
        print(response.json()['result_message'])
        # 输入验证码进行登入
        code = str(input('请输入验证码:'))
        url_code = 'https://kyfw.12306.cn/passport/web/login'
        # 构造表单数据
        data_code = {
            'sessionId':'',
            'sig':'',
            'if_check_slide_passcode_token':'',
            'scene':'',
            'checkMode':'0',
            'randCode':code,                              # 验证码
            'username':username,                          # 账号
            'password':mima_SM4(password),                # 密码,需要进行js逆向
            'appid':'otn',
        }
        # 进行登入
        response = session.post(url_code,data=data_code)
        # print(session.cookies)
        if '成功' in response.json()['result_message']:
            print(response.json()['result_message'])          # 输出登入成功

            # 输出登入者信息
            response = session.post('https://kyfw.12306.cn/otn/login/conf')
            print(response)
            return session
        else:
            print(response.json()['result_message'])
            return None
    else:
        print(response.json()['result_message'])
        return None

if __name__ == '__main__':
    username = str(input('请输入手机号：'))
    castNum = str(input('请输入身份证后4位：'))
    password = str(input('请输入密码：'))
    print('正在登入...')
    USER = login(username,castNum,password)